{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# rStar-Math: Linear Algebra Examples\n",
    "\n",
    "This notebook demonstrates how rStar-Math handles linear algebra problems with visualizations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from src.core.mcts import MCTS\n",
    "from src.core.ppm import ProcessPreferenceModel\n",
    "from src.models.model_interface import ModelFactory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "# Initialize components\n",
    "mcts = MCTS.from_config_file('config/default.json')\n",
    "ppm = ProcessPreferenceModel.from_config_file('config/default.json')\n",
    "model = ModelFactory.create_model('openai', os.getenv('OPENAI_API_KEY'), 'config/default.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Matrix Operations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "matrix_problems = [\n",
    "    \"Find the determinant of [[1, 2], [3, 4]]\",\n",
    "    \"Solve the system of equations: 2x + y = 5, x - y = 1\",\n",
    "    \"Find the eigenvalues of [[2, 1], [1, 2]]\",\n",
    "    \"Calculate the inverse of [[1, 2], [3, 4]]\"\n",
    "]\n",
    "\n",
    "def visualize_matrix(matrix_str: str):\n",
    "    \"\"\"Visualize matrix as a heatmap.\"\"\"\n",
    "    matrix = np.array(eval(matrix_str))\n",
    "    plt.figure(figsize=(8, 6))\n",
    "    plt.imshow(matrix, cmap='viridis')\n",
    "    plt.colorbar()\n",
    "    for i in range(matrix.shape[0]):\n",
    "        for j in range(matrix.shape[1]):\n",
    "            plt.text(j, i, f'{matrix[i,j]:.2f}', ha='center', va='center')\n",
    "    plt.title('Matrix Visualization')\n",
    "    plt.show()\n",
    "\n",
    "for problem in matrix_problems:\n",
    "    print(f\"Problem: {problem}\\n\")\n",
    "    action, trajectory = mcts.search(problem)\n",
    "    \n",
    "    print(\"Solution Steps:\")\n",
    "    for step in trajectory:\n",
    "        confidence = ppm.evaluate_step(step['state'], model)\n",
    "        print(f\"- {step['state']}\")\n",
    "        print(f\"  Confidence: {confidence:.2f}\\n\")\n",
    "        \n",
    "    # Visualize matrix if present in problem\n",
    "    if '[[' in problem:\n",
    "        matrix_str = problem[problem.find('[['):problem.find(']]')+2]\n",
    "        visualize_matrix(matrix_str)\n",
    "    print(\"-\" * 50 + \"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Vector Spaces and Transformations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "def plot_vector_transformation(matrix: np.ndarray):\n",
    "    \"\"\"Visualize linear transformation.\"\"\"\n",
    "    fig = plt.figure(figsize=(12, 5))\n",
    "    \n",
    "    # Original vectors\n",
    "    ax1 = fig.add_subplot(121)\n",
    "    vectors = np.array([[1, 0], [0, 1]])\n",
    "    ax1.quiver([0, 0], [0, 0], vectors[:, 0], vectors[:, 1],\n",
    "               angles='xy', scale_units='xy', scale=1)\n",
    "    ax1.set_xlim(-2, 2)\n",
    "    ax1.set_ylim(-2, 2)\n",
    "    ax1.grid(True)\n",
    "    ax1.set_title('Original Vectors')\n",
    "    \n",
    "    # Transformed vectors\n",
    "    ax2 = fig.add_subplot(122)\n",
    "    transformed = np.dot(vectors, matrix)\n",
    "    ax2.quiver([0, 0], [0, 0], transformed[:, 0], transformed[:, 1],\n",
    "               angles='xy', scale_units='xy', scale=1)\n",
    "    ax2.set_xlim(-2, 2)\n",
    "    ax2.set_ylim(-2, 2)\n",
    "    ax2.grid(True)\n",
    "    ax2.set_title('Transformed Vectors')\n",
    "    \n",
    "    plt.show()\n",
    "\n",
    "# Example transformations\n",
    "transformations = [\n",
    "    np.array([[2, 0], [0, 2]]),  # Scaling\n",
    "    np.array([[0, -1], [1, 0]]),  # Rotation\n",
    "    np.array([[1, 1], [0, 1]])   # Shear\n",
    "]\n",
    "\n",
    "for matrix in transformations:\n",
    "    print(f\"Transformation Matrix:\\n{matrix}\\n\")\n",
    "    plot_vector_transformation(matrix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Eigenvalues and Eigenvectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "def plot_eigenvectors(matrix: np.ndarray):\n",
    "    \"\"\"Visualize eigenvectors and their transformations.\"\"\"\n",
    "    eigenvals, eigenvecs = np.linalg.eig(matrix)\n",
    "    \n",
    "    plt.figure(figsize=(8, 8))\n",
    "    \n",
    "    # Plot original vectors\n",
    "    for i, vec in enumerate(eigenvecs.T):\n",
    "        plt.quiver(0, 0, vec[0], vec[1], angles='xy', scale_units='xy',\n",
    "                  scale=1, color='blue', label=f'Eigenvector {i+1}')\n",
    "        \n",
    "    # Plot transformed vectors\n",
    "    transformed = np.dot(matrix, eigenvecs)\n",
    "    for i, vec in enumerate(transformed.T):\n",
    "        plt.quiver(0, 0, vec[0], vec[1], angles='xy', scale_units='xy',\n",
    "                  scale=1, color='red', label=f'Transformed {i+1}')\n",
    "    \n",
    "    plt.xlim(-2, 2)\n",
    "    plt.ylim(-2, 2)\n",
    "    plt.grid(True)\n",
    "    plt.legend()\n",
    "    plt.title('Eigenvectors and Their Transformations')\n",
    "    plt.show()\n",
    "    \n",
    "    print(\"Eigenvalues:\")\n",
    "    for i, val in enumerate(eigenvals):\n",
    "        print(f\"λ{i+1} = {val:.2f}\")\n",
    "\n",
    "# Example matrices\n",
    "matrices = [\n",
    "    np.array([[2, 1], [1, 2]]),  # Symmetric matrix\n",
    "    np.array([[0, -1], [1, 0]]),  # Rotation matrix\n",
    "    np.array([[3, 1], [0, 2]])   # Upper triangular matrix\n",
    "]\n",
    "\n",
    "for matrix in matrices:\n",
    "    print(f\"\\nMatrix:\\n{matrix}\")\n",
    "    plot_eigenvectors(matrix)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
